/**
 * Copyright (C) 2007 The AsyncMail Group <asyncmail@googlegroups.com>
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *         http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.googlecode.asyncmail.smtpserver.handler.fastfail;

import java.util.Collection;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import javax.sql.DataSource;

import org.apache.commons.logging.Log;
import org.springframework.beans.factory.annotation.Required;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.simple.ParameterizedRowMapper;
import org.springframework.jdbc.core.simple.SimpleJdbcTemplate;

import com.googlecode.asyncmail.common.DSNStatus;
import com.googlecode.asyncmail.common.MailAddress;
import com.googlecode.asyncmail.dnsservice.DNSService;
import com.googlecode.asyncmail.dnsservice.NetMatcher;
import com.googlecode.asyncmail.smtp.SMTPRetCode;
import com.googlecode.asyncmail.smtpserver.SMTPSession;
import com.googlecode.asyncmail.smtpserver.hook.HookResult;
import com.googlecode.asyncmail.smtpserver.hook.HookReturnCode;
import com.googlecode.asyncmail.smtpserver.hook.RcptHook;
import com.googlecode.asyncmail.util.TimeConverter;

/**
 * GreylistHandler which can be used to activate Greylisting
 */
public class GreylistHandler implements RcptHook {

    // 1 hour
    private long tempBlockTime = 3600000;

    // 36 days
    private long autoWhiteListLifeTime = 3110400000L;

    // 4 hours
    private long unseenLifeTime = 14400000;

    // TODO Remove hardcoding . Maybe use ibatis.apache.org ?
    private String selectQuery = "SELECT create_time,connectioncount FROM greylist WHERE ipaddress = ? AND sender = ? AND recip = ?";

    private String insertQuery = "INSERT INTO greylist (ipaddress,sender,recip,connectioncount,create_time) values (?,?,?,?,?)";

    private String deleteQuery = "DELETE FROM greylist WHERE create_time < ? AND connectioncount = 0";

    private String deleteAutoWhiteListQuery = "DELETE FROM greylist WHERE create_time < ?";

    private String updateQuery = "UPDATE greylist SET create_time = ? , connectioncount = ? WHERE ipaddress = ? AND sender = ? AND recip = ?";

    private String createTableQuery = "CREATE TABLE greylist (ipaddress varchar (20) NOT NULL,sender varchar (255) NOT NULL, recip varchar (255) NOT NULL, connectioncount int NOT NULL, create_time timestamp NOT NULL, PRIMARY KEY (ipaddress,sender,recip))";

    private String greyListTableName = "greylist";

    private NetMatcher wNetworks;

    private DataSource dataSource;

    private SimpleJdbcTemplate simpleJdbcTemplate;

    private Log logger;

    private DNSService dnsService;

    @Required
    public void setDataSource(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    @Required
    public void setLogger(Log logger) {
        this.logger = logger;
    }

    @Required
    public void setDnsService(DNSService dnsService) {
        this.dnsService = dnsService;
    }

    protected Log getLogger() {
        return logger;
    }

    public void setWhitelistedNetworks(Collection<String> nets) {
        wNetworks = new NetMatcher(nets, dnsService, logger);
    }

    /**
     * Init the handler and create the table in database if not exist
     */
    public void initialize() throws Exception {

        ResultSet result = dataSource.getConnection().getMetaData().getTables(
                null, null, greyListTableName, null);
        simpleJdbcTemplate = new SimpleJdbcTemplate(dataSource);

        // Check if table exists. If not create it
        if (result.next() == false) {
            createTable();
        }

    }

    /**
     * Setup the temporary blocking time
     * 
     * @param tempBlockTime
     *                The temporary blocking time
     */
    public void setTempBlockTime(String tempBlockTime) {
        setTempBlockTime(TimeConverter.getMilliSeconds(tempBlockTime));
    }

    public void setTempBlockTime(long tempBlockTime) {
        this.tempBlockTime = tempBlockTime;
    }

    /**
     * Setup the autowhitelist lifetime for which we should whitelist a triplet.
     * After this lifetime the record will be deleted
     * 
     * @param autoWhiteListLifeTime
     *                The lifeTime
     */
    public void setAutoWhiteListLifeTime(String autoWhiteListLifeTime) {
        setAutoWhiteListLifeTime(TimeConverter
                .getMilliSeconds(autoWhiteListLifeTime));
    }

    public void setAutoWhiteListLifeTime(long autoWhiteListLifeTime) {
        this.autoWhiteListLifeTime = autoWhiteListLifeTime;
    }

    /**
     * Set up the liftime of only once seen triplet. After this liftime the
     * record will be deleted
     * 
     * @param unseenLifeTime
     *                The lifetime
     */
    public void setUnseenLifeTime(String unseenLifeTime) {
        setUnseenLifeTime(TimeConverter.getMilliSeconds(unseenLifeTime));
    }

    public void setUnseenLifeTime(long unseenLifeTime) {
        this.unseenLifeTime = unseenLifeTime;
    }

    private HookResult doGreyListCheck(SMTPSession session,
            MailAddress senderAddress, MailAddress recipAddress) {
        String recip = "";
        String sender = "";

        if (recipAddress != null)
            recip = recipAddress.toString();
        if (senderAddress != null)
            sender = senderAddress.toString();

        long time = System.currentTimeMillis();
        String ipAddress = session.getRemoteIPAddress();

        try {
            long createTimeStamp = 0;
            int count = 0;

            // get the timestamp when he triplet was last seen
            GreyListData data = getGreyListData(ipAddress, sender, recip);

            createTimeStamp = data.getTimeStamp();
            count = data.getCount();

            getLogger().debug(
                    "Triplet " + ipAddress + " | " + sender + " | " + recip
                            + " -> TimeStamp: " + createTimeStamp);

            // if the timestamp is bigger as 0 we have allready a triplet stored
            if (createTimeStamp > 0) {
                long acceptTime = createTimeStamp + tempBlockTime;

                if ((time < acceptTime) && (count == 0)) {
                    return new HookResult(
                            HookReturnCode.DENYSOFT,
                            SMTPRetCode.LOCAL_ERROR,
                            DSNStatus.getStatus(DSNStatus.TRANSIENT,
                                    DSNStatus.NETWORK_DIR_SERVER)
                                    + " Temporary rejected: Reconnect to fast. Please try again later");
                } else {

                    getLogger().debug(
                            "Update triplet " + ipAddress + " | " + sender
                                    + " | " + recip + " -> timestamp: " + time);

                    // update the triplet..
                    updateTriplet(ipAddress, sender, recip, count, time);

                }
            } else {
                getLogger().debug(
                        "New triplet " + ipAddress + " | " + sender + " | "
                                + recip);

                // insert a new triplet
                insertTriplet(ipAddress, sender, recip, count, time);

                // Tempory block on new triplet!
                return new HookResult(HookReturnCode.DENYSOFT,
                        SMTPRetCode.LOCAL_ERROR, DSNStatus.getStatus(
                                DSNStatus.TRANSIENT,
                                DSNStatus.NETWORK_DIR_SERVER)
                                + " Temporary rejected: Please try again later");
            }

            // some kind of random cleanup process
            if (Math.random() > 0.99) {
                // cleanup old entries

                getLogger().debug("Delete old entries");

                cleanupAutoWhiteListGreyList(time - autoWhiteListLifeTime);
                cleanupGreyList(time - unseenLifeTime);
            }

        } catch (DataAccessException e) {
            // just log the exception
            getLogger().error("Error on SQLquery: " + e.getMessage());
        }
        return new HookResult(HookReturnCode.DECLINED);
    }

    /**
     * Return the GreyListData object for the given input
     * 
     * @param ipAddress
     * @param sender
     * @param recip
     * @return
     */
    private GreyListData getGreyListData(String ipAddress, String sender,
            String recip) {
        try {
            return simpleJdbcTemplate.queryForObject(selectQuery,
                    new GreyListResultSetRawMapper(), new Object[] { ipAddress,
                            sender, recip });
        } catch (DataAccessException e) {
            logger.debug("No entries in table");
            return new GreyListData();
        }

    }

    /**
     * Insert new triplet in the store
     * 
     * @param ipAddress
     * @param sender
     * @param recip
     * @param count
     * @param createTime
     * @throws DataAccessException
     */
    private void insertTriplet(String ipAddress, String sender, String recip,
            int count, long createTime) throws DataAccessException {

        simpleJdbcTemplate.update(insertQuery, new Object[] { ipAddress,
                sender, recip, count, new Timestamp(createTime) });
    }

    /**
     * Update the triplet
     * 
     * @param ipAddress
     * @param sender
     * @param recip
     * @param count
     * @param time
     * @throws DataAccessException
     */
    private void updateTriplet(String ipAddress, String sender, String recip,
            int count, long time) throws DataAccessException {
        simpleJdbcTemplate.update(updateQuery, new Object[] {
                new Timestamp(time), (count + 1), ipAddress, sender, recip });
    }

    /**
     * Cleanup the autowhitelist
     * 
     * @param time
     * @throws DataAccessException
     */
    private void cleanupAutoWhiteListGreyList(long time)
            throws DataAccessException {
        simpleJdbcTemplate.update(deleteAutoWhiteListQuery,
                new Object[] { new Timestamp(time) });
    }

    /**
     * Cleanup the greylist
     * 
     * @param time
     */
    private void cleanupGreyList(long time) throws DataAccessException {
        simpleJdbcTemplate.update(deleteQuery, new Object[] { new Timestamp(
                time) });
    }

    /**
     * Create the table if not exists.
     */
    private void createTable() throws DataAccessException {
        simpleJdbcTemplate.update(createTableQuery, new Object[] {});
        StringBuffer logBuffer = new StringBuffer(64).append("Created table '")
                .append(greyListTableName);
        getLogger().info(logBuffer.toString());
    }

    /**
     * @see com.googlecode.asyncmail.smtpserver.hook.RcptHook#doRcpt(com.googlecode.asyncmail.smtpserver.SMTPSession,
     *      org.apache.mailet.MailAddress, org.apache.mailet.MailAddress)
     */
    public HookResult doRcpt(SMTPSession session, MailAddress sender,
            MailAddress rcpt) {
        if ((wNetworks == null)
                || (!wNetworks.matchInetNetwork(session.getRemoteIPAddress()))) {
            return doGreyListCheck(session, sender, rcpt);
        } else {
            getLogger().info(
                    "IpAddress " + session.getRemoteIPAddress()
                            + " is whitelisted. Skip greylisting.");
        }

        return new HookResult(HookReturnCode.DECLINED);
    }

    /**
     * Private class to represent the GreyListData
     */
    private class GreyListData {
        private long timeStamp = 0;

        private int count = 0;

        public GreyListData() {
        }

        public GreyListData(long timeStamp, int count) {
            this.timeStamp = timeStamp;
            this.count = count;
        }

        public long getTimeStamp() {
            return timeStamp;
        }

        public int getCount() {
            return count;
        }
    }

    /**
     * Greylist implementation of ParameterizedRowMapper
     * 
     */
    private class GreyListResultSetRawMapper implements
            ParameterizedRowMapper<GreyListData> {

        /**
         * @see org.springframework.jdbc.core.simple.ParameterizedRowMapper#mapRow(java.sql.ResultSet,
         *      int)
         */
        public GreyListData mapRow(ResultSet arg0, int arg1)
                throws SQLException {
            return new GreyListData(arg0.getTimestamp(1).getTime(), arg0
                    .getInt(2));
        }

    }
}
